execfile('v4/base.py')

class LabelF(Base):
    def __init__(self):
        self.bar_count_default = 50
        self.end_date_str = '20250831'  # 如果需要可以取消注释
        # self.start_date_str = '2025-06-16'  # 如果需要可以取消注释

    """ 标注后续表现（计算未来N日收益率）"""
    def label_future(self,df, days):
        col = 'ret'+str(days)
        df['future'] = df.groupby('code')['close'].shift(-days)
        df[col] = (df['future'] / df['close'] - 1)
        # df = df[['date','code',col]]
        return df
    
    def mark_label(self,df):
        df = df.sort_values(['code', 'date'])
        #计算多个时间窗口的收益率
        label_array = [1,2,3,5]
        for days in label_array:
            df = self.label_future(df, days)
        # 只保留需要的列
        keep_cols = ['date', 'code'] + [f'ret{d}' for d in label_array]
        df = df[keep_cols]
        df['date'] = pd.to_datetime(df['date'])
        df = df.round(2)
        return df

    def binarize_ret_10cm_gt4(self,df):
        # 初始化所有标签为0（默认值）
        df['label'] = 0
        # 优先级1: 满足任意一个高增长条件 → label=3
        cond3 = ((df['ret1'] > 0.06) & (df['ret2'] > 0.12) | (df['ret3'] > 0.18) )
        df.loc[cond3, 'label'] = 3
        # 优先级2: 满足任意一个中等增长条件 → label=2 (且未标记为3)
        cond2 = (df['ret1'] > 0.03) | (df['ret2'] > 0.06) |  (df['ret3'] > 0.12)
        df.loc[cond2 & (df['label'] == 0), 'label'] = 2  # 仅更新未标记的样本
        # 优先级3: 满足任意一个低增长条件 → label=1 (且未标记为2或3)
        cond1 = (df['ret1'] > 0) | (df['ret2'] > 0) |  (df['ret3'] > 0)
        df.loc[cond1 & (df['label'] == 0), 'label'] = 1  # 仅更新未标记的样本
        # 优先级4: 所有收益率都小于0 → 保持label=0
        # 不需要额外操作，因为默认已经是0
        return df
        
    def binarize_ret_10cm_zt(self,df):
        # 初始化所有标签为0（默认值）
        df['label'] = 0
        # 优先级1: 满足任意一个高增长条件 → label=3
        cond3 = (df['ret2'] > 0.12) | (df['ret3'] > 0.21)  
        df.loc[cond3, 'label'] = 3
        # 优先级2: 满足任意一个中等增长条件 → label=2 (且未标记为3)
        cond2 = (df['ret1'] > 0.04) | (df['ret2'] > 0.05) |  (df['ret3'] > 0.10)
        df.loc[cond2 & (df['label'] == 0), 'label'] = 2  # 仅更新未标记的样本
        # 优先级3: 满足任意一个低增长条件 → label=1 (且未标记为2或3)
        cond1 = (df['ret1'] > 0) | (df['ret2'] > 0) |  (df['ret3'] > 0)
        df.loc[cond1 & (df['label'] == 0), 'label'] = 1  # 仅更新未标记的样本
        # 优先级4: 所有收益率都小于0 → 保持label=0
        # 不需要额外操作，因为默认已经是0
        return df

    def binarize_outlier(self,df):
        # 初始化所有标签为0（默认值）
        df['label'] = 0
        # 优先级1: 满足任意一个高增长条件 → label=3
        cond3 = (df['ret5'] > 0.40)   
        df.loc[cond3, 'label'] = 3
        # 优先级2: 满足任意一个中等增长条件 → label=2 (且未标记为3)
        cond2 = (df['ret5'] > 0.2) 
        df.loc[cond2 & (df['label'] == 0), 'label'] = 2  # 仅更新未标记的样本
        # 优先级3: 满足任意一个低增长条件 → label=1 (且未标记为2或3)
        cond1 = (df['ret5'] > 0) 
        df.loc[cond1 & (df['label'] == 0), 'label'] = 1  # 仅更新未标记的样本
        # 优先级4: 所有收益率都小于0 → 保持label=0
        # 不需要额外操作，因为默认已经是0
        return df